#!/usr/bin/bash
# SPDX-License-Identifier: GPL-2.0-only
# Copyright (c) 2023 Intel Corporation

# Author:   Hongyu Ning <hongyu.ning@intel.com>
#
# History:  16, Oct., 2023 - Hongyu Ning - creation


# @desc This script do basic TD attestation check in TDX Guest VM
#       test binary is based on kselftest linux/tools/testing/selftests/tdx implementation

###################### Variables ######################
SCRIPT_DIR="$( cd "$( dirname "$0" )" && pwd )"
echo "$SCRIPT_DIR"
source common.sh

# TSM based attest sysfs
tsm_config=/sys/kernel/config/tsm/report

while getopts :t: arg; do
  case $arg in
    t)
      ATTEST_CASE=$OPTARG
      ;;
    *)
      test_print_err "Must supply an argument to -$OPTARG."
      exit 1
      ;;
  esac
done

###################### Functions ######################
attest_result() {
  selftest_item=$1
  case "$selftest_item" in
    global.verify_report)
      selftest_num=1
      ;;
    global.verify_reportmac)
      selftest_num=2
      ;;
    global.verify_rtmr_extend)
      selftest_num=3
      ;;
    global.verify_quote)
      selftest_num=4
      ;;
  esac
  test_print_trc "TD attestation - $selftest_item start."
  if [ -f "attest.log" ]; then
    rm -rf attest.log
  fi
  ./tdx_guest_test | tee attest.log
  results=$(grep "not ok $selftest_num $selftest_item" attest.log)
  if [ -z "$results" ]; then
    test_print_trc "TD attestation - $selftest_item PASS."
  else
    die "TD attestation - $selftest_item FAIL."
    return 1
  fi
}

tsm_get_quote_pre_check() {
  report="$tsm_config"/$1
  mkdir "$report"
  # tsm td quote sysfs precheck
  if [[ -f "$report"/generation ]] && \
    [[ $(cat "$report"/generation) -eq 0 ]]; then
    test_print_trc "TD TSM attest quote generation pre-check correct"
  else
    die "TD TSM attest quote generation not exists or not equal to 0"
    return 1
  fi
  if [[ -f "$report"/provider ]] && \
    [[ $(cat "$report"/provider) == "tdx_guest" ]]; then
    test_print_trc "TD TSM attest quote provider pre-check correct"
  else
    die "TD TSM attest quote provider not exists or not equal to 'tdx_guest'"
    return 1
  fi
  if [[ -f "$report"/inblob ]]; then
    test_print_trc "TD TSM attest quote inblob pre-check correct"
  else
    die "TD TSM attest quote inblob not exists"
    return 1
  fi
  if [[ -f "$report"/outblob ]]; then
    test_print_trc "TD TSM attest quote outblob pre-check correct"
  else
    die "TD TSM attest quote outblob not exists"
    return 1
  fi
  rmdir "$report"
}

tsm_get_quote() {
  report="$tsm_config"/$1
  mkdir "$report"
  # generate quote check
  generation_before=$(cat "$report"/generation)
  test_print_trc "tsm quote generation before trigger inblob: $generation_before"
  if ! dd if=/dev/urandom bs=64 count=1 > "$report"/inblob; then
    die "TD TSM attest quote inblob write failed"
    return 1
  else
    test_print_trc "TD TSM attest quote triggered success once"
  fi
  generation_after=$(cat "$report"/generation)
  test_print_trc "tsm quote generation after trigger inblob: $generation_after"
  if [[ "$generation_after" -gt "$generation_before" ]]; then
    test_print_trc "TD TSM attest quote inblob triggered 1 more"
  else
    die "TD TSM attest quote inblob triggered generation failed"
    return 1
  fi
  quote_return=$(hexdump -C "$report"/outblob 2>&1 >/dev/null)
  test_print_trc "hexdump -C outblob return: $quote_return"
  if echo "$quote_return" | grep "Connection timed out"; then
    test_print_trc "TD TSM attest quote generated from @inblob success"
  else
    die "TD TSM attest quote generated by inblob failed"
    return 1
  fi
  rmdir "$report"
}

tsm_get_quote_negative() {
  report="$tsm_config"/$1
  mkdir "$report"
  # negative generate quote check
  generation_before=$(cat "$report"/generation)
  test_print_trc "tsm quote generation before trigger inblob: $generation_before"
  if ! dd if=/dev/urandom bs=32 count=1 > "$report"/inblob; then
    die "TD TSM attest quote inblob write failed"
    return 1
  else
    test_print_trc "TD TSM attest quote triggered success once"
  fi
  generation_after=$(cat "$report"/generation)
  test_print_trc "tsm quote generation after trigger inblob: $generation_after"
  if [[ "$generation_after" -gt "$generation_before" ]]; then
    test_print_trc "TD TSM attest quote inblob triggered 1 more"
  else
    die "TD TSM attest quote inblob triggered generation failed"
    return 1
  fi
  quote_return=$(hexdump -C "$report"/outblob 2>&1 >/dev/null)
  test_print_trc "hexdump -C outblob return: $quote_return"
  if echo "$quote_return" | grep "Invalid argument"; then
    test_print_trc "TD TSM attest quote negative generated from @inblob success"
  else
    die "TD TSM attest quote negative generated by inblob failed"
    return 1
  fi
  rmdir "$report"
}

###################### Do Works ######################

case "$ATTEST_CASE" in
  global.verify_report)
    attest_result "$ATTEST_CASE"
    ;;
  global.verify_reportmac)
    attest_result "$ATTEST_CASE"
    ;;
  global.verify_rtmr_extend)
    attest_result "$ATTEST_CASE"
    ;;
  global.verify_quote)
    attest_result "$ATTEST_CASE"
    ;;
  tsm.get_quote.precheck)
    tsm_get_quote_pre_check report0
    ;;
  tsm.get_quote)
    tsm_get_quote report1
    ;;
  tsm.get_quote.negative)
    tsm_get_quote_negative report2
    ;;
  :)
    test_print_err "Must specify the attest case option by [-t]"
    exit 1
    ;;
  \?)
    test_print_err "Input test case option $ATTEST_CASE is not supported"
    exit 1
    ;;
esac